package com.sh.data.engine.domain.shims.hdfs.util;

import com.google.common.base.Joiner;
import com.google.common.collect.Lists;
import lombok.Builder;
import lombok.Data;
import lombok.extern.slf4j.Slf4j;
import org.apache.commons.collections.CollectionUtils;
import org.apache.commons.lang3.StringUtils;
import org.apache.hadoop.conf.Configuration;
import org.apache.hadoop.fs.FSDataOutputStream;
import org.apache.hadoop.fs.FileSystem;
import org.apache.hadoop.fs.Path;
import org.buildobjects.process.ProcBuilder;
import org.buildobjects.process.ProcResult;

import java.io.ByteArrayOutputStream;
import java.io.IOException;
import java.io.OutputStream;
import java.net.URI;
import java.nio.charset.StandardCharsets;
import java.util.Arrays;
import java.util.Collections;
import java.util.LinkedList;
import java.util.List;

/**
 * @author: mengzheng.mhc
 * @date: 2024/6/11 16:57
 */
@Slf4j
public class HdfsUtil {

    private static final String HIVE_WAREHOUSE_PATH = "/user/hive/warehouse/";

    private static final String DEFAULT_DB = "default.db";

    private static final Joiner SPACE_JOINER = Joiner.on(" ");

    private static final String HDFS_CMD = "/usr/bin/hdfs";

    public static void testConnection(String user, String dsLink, String hdfsConfig)
        throws IOException {

        if (!dsLink.startsWith("hdfs://")) {
            log.error("url格式错误，连接失败");
        }
        FileSystem hdfs = getFileSystem(null, user, null, dsLink, null, null, null);
        hdfs.listFiles(new Path("/"), false);
    }

    public static FileSystem getFileSystem(
        String krb5Conf,
        String username,
        String keyTabPath,
        String hdfsUrl,
        String nameservices,
        String namenodesAddrStr,
        String namenodesStr) {

        try {

            FileSystem hdfs = null;
            URI uri = new URI(hdfsUrl);
            if (StringUtils.isBlank(nameservices)) {
                Configuration conf = new Configuration();
                hdfs = FileSystem.get(uri, conf, username);
            } else {
                Configuration conf = new Configuration(false);
                String[] namenodesAddr = namenodesAddrStr.split(",");
                String[] namenodes = namenodesStr.split(",");
                conf.set("fs.defaultFS", "hdfs://" + nameservices);
                conf.set("dfs.nameservices", nameservices);
                conf.set("dfs.ha.namenodes." + nameservices, namenodesStr);
                for (int i = 0; i < namenodes.length; i++) {
                    conf.set(
                        "dfs.namenode.rpc-address." + nameservices + "." + namenodes[i], namenodesAddr[i]);
                }
                conf.set(
                    "dfs.client.failover.proxy.provider." + nameservices,
                    "org.apache.hadoop.hdfs.server.namenode.ha.ConfiguredFailoverProxyProvider");
                hdfs = FileSystem.get(uri, conf, username);
            }

            return hdfs;
        } catch (Exception e) {
            log.error("", e);
        }

        return null;
    }

    public static void closeFSDataOutputStream(FSDataOutputStream out) {
        if (out != null) {
            try {
                out.close();
            } catch (IOException e) {
                log.error("", e);
            }
        }
    }

    public static void closeFileSystem(FileSystem hdfs) {
        if (hdfs != null) {
            try {
                hdfs.close();
            } catch (IOException e) {
                log.error("", e);
            }
        }
    }

    public static List<TableInfo> getHiveTableSize() {
        String path = HIVE_WAREHOUSE_PATH + "*db";
        List<String> args = new LinkedList<String>();
        args.add("-");
        args.add("hdfs");
        args.add("-s");
        args.add("/bin/bash");
        args.add("-c");
        args.add(String.format("hdfs dfs -du %s", path));
        ByteArrayOutputStream out = new ByteArrayOutputStream();
        exeWithLog(out, args.toArray(new String[args.size()]));
        String raw = new String(out.toByteArray(), StandardCharsets.UTF_8);
        return parseHiveTableInfo(raw);
    }

    private static List<TableInfo> parseHiveTableInfo(String raw) {
        if (StringUtils.isBlank(raw)) {
            return Collections.emptyList();
        }

        String[] array = StringUtils.split(raw, "\n");
        if (array == null || array.length == 0) {
            return Collections.emptyList();
        }

        List<TableInfo> list = new LinkedList<>();
        for (String line : array) {
            if (StringUtils.isBlank(line)) {
                continue;
            }
            line = StringUtils.trim(line);
            String[] values = StringUtils.split(line, " ");
            if (values == null || values.length != 3) {
                continue;
            }
            Long size = parseLong(values[0]);
            String dbName = parseDbName(values[2]);
            String tableName = parseTableName(values[2], dbName);
            if (size == null || dbName == null || tableName == null) {
                continue;
            }
            list.add(TableInfo.builder().dbName(dbName).tableName(tableName).size(size).build());
        }
        return list;
    }

    public static Long parseLong(String str) {
        if (StringUtils.isBlank(str)) {
            return null;
        }
        str = StringUtils.trim(str);
        try {
            return Long.parseLong(str);
        } catch (NumberFormatException e) {
        }
        return null;
    }

    private static String parseDbName(String str) {
        if (StringUtils.isBlank(str)) {
            return null;
        }

        str = StringUtils.trim(str);
        String db = StringUtils.substringBetween(str, HIVE_WAREHOUSE_PATH, DEFAULT_DB) + DEFAULT_DB;
        db = StringUtils.substringBeforeLast(db, ".");
        return db;
    }

    private static String parseTableName(String str, String dbName) {
        if (StringUtils.isBlank(str)) {
            return null;
        }
        str = StringUtils.trim(str);
        return StringUtils.substringAfterLast(str, "/");
    }

    public static boolean exeWithLog(OutputStream out, String... args) {
        return exeWithLog(out, null, args);
    }

    public static boolean exeWithLog(OutputStream out, OutputStream err, String... args) {
        try {
            ProcBuilder builder = new ProcBuilder("/usr/bin/su").withArgs(args).withNoTimeout();
            if (out != null) {
                builder.withOutputStream(out);
            }
            if (err != null) {
                builder.withErrorStream(err);
            }
            ProcResult result = builder.run();
            if (result.getExitValue() == 0) {
                return true;
            }
        } catch (Exception e) {
        }
        return false;
    }

    public static boolean uploadFile(String user, String sourceFilePath, String targetDir)
        throws IOException {
        String[] mkdirsArgs = {HDFS_CMD, "dfs", "-mkdir", "-p", targetDir};

        String[] uploadArgs = {HDFS_CMD, "dfs", "-put", "-f", sourceFilePath, targetDir};

        try {
            runCmd(user, mkdirsArgs);
            runCmd(user, uploadArgs);
            return true;
        } catch (IOException e) {
            throw e;
        }
    }

    public static void deleteDir(String user, List<String> dirList) throws IOException {
        if (CollectionUtils.isEmpty(dirList)) {
            return;
        }

        List<String> argList = Lists.newArrayList(HDFS_CMD, "dfs", "-rm", "-r", "-skipTrash");
        argList.addAll(dirList);

        String[] args = new String[argList.size()];
        argList.toArray(args);

        try {
            runCmd(user, args);
        } catch (IOException e) {
            throw e;
        }
    }

    private static int runCmd(String user, String... args) throws IOException {
        String cmd = "/usr/bin/su";
        String[] cmdArgs = new String[]{"-", user, "-s", "/bin/bash", "-c", SPACE_JOINER.join(args)};
        try (ByteArrayOutputStream out = new ByteArrayOutputStream()) {
            ProcBuilder builder = new ProcBuilder(cmd).withArgs(cmdArgs).withNoTimeout();
            builder.withOutputStream(out);
            ProcResult result = builder.run();
            return result.getExitValue();
        } catch (IOException e) {
            log.error("execute command {} {} error", cmd, Arrays.toString(cmdArgs), e);
            throw e;
        }
    }

    @Builder
    @Data
    public static class TableInfo {
        private String dbName;
        private String tableName;
        private Long size;
    }
}
